package com.zzk.audiovideodemo.extractorMuxer;

import android.content.Context;
import android.media.MediaCodec;
import android.media.MediaExtractor;
import android.media.MediaFormat;
import android.media.MediaMuxer;
import android.text.TextUtils;
import android.widget.Toast;

import com.blankj.utilcode.util.LogUtils;
import com.blankj.utilcode.util.ToastUtils;

import java.io.File;
import java.io.IOException;
import java.nio.ByteBuffer;

/**
 * @ProjectName: AudioVideoDemo
 * @Package: com.zzk.audiovideodemo.extractorMuxer
 * @ClassName: ExtractorMuxerUtil
 * @Description:
 * @Author: brilliantzhao
 * @CreateDate: 2021.3.24 14:32
 * @UpdateUser:
 * @UpdateDate: 2021.3.24 14:32
 * @UpdateRemark:
 * @Version: 1.0.0
 */

public class ExtractorMuxerUtil {

    private Context mContext;

    public ExtractorMuxerUtil(Context context) {
        this.mContext = context;
    }

    /**
     * 提取mp4中的音轨和视频轨道
     */
    public void extractorAndMuxerMP4() {
        File inputFile = new File(mContext.getExternalCacheDir().getAbsolutePath() + "/MediaRecord/video_record.mp4");
        if (!inputFile.exists()) {
            Toast.makeText(mContext, "mp4文件不存在", Toast.LENGTH_SHORT).show();
            return;
        }
        //数据提取(解封装)
        //1. 构造MediaExtractor
        MediaExtractor mediaExtractor = new MediaExtractor();
        try {
            //2.设置数据源
            mediaExtractor.setDataSource(inputFile.getAbsolutePath());
            //3. 获取轨道数
            int trackCount = mediaExtractor.getTrackCount();
            LogUtils.i("demuxerMP4: trackCount=" + trackCount);
            //遍历轨道，查看音频轨或者视频轨道信息
            for (int i = 0; i < trackCount; i++) {
                //4. 获取某一轨道的媒体格式
                MediaFormat trackFormat = mediaExtractor.getTrackFormat(i);
                String keyMime = trackFormat.getString(MediaFormat.KEY_MIME);
                LogUtils.i("demuxerMp4: keyMime=" + keyMime);
                if (TextUtils.isEmpty(keyMime)) {
                    continue;
                }
                //5.通过mime信息识别音轨或视频轨道，打印相关信息
                if (keyMime.startsWith("video/")) {
                    File outputFile = extractorAndMuxer(mContext, mediaExtractor, i, "/single_video.mp4");
                    LogUtils.i("纯视频文件路径：" + outputFile.getAbsolutePath());
                    ToastUtils.showLong("纯视频文件路径：" + outputFile.getAbsolutePath());
                    LogUtils.i("extractorAndMuxerMP4: videoWidth=" + trackFormat.getInteger(MediaFormat.KEY_WIDTH) + " videoHeight=" + trackFormat.getInteger(MediaFormat.KEY_HEIGHT));

                } else if (keyMime.startsWith("audio/")) {
                    File outputFile = extractorAndMuxer(mContext, mediaExtractor, i, "/single_audio.aac");
                    LogUtils.i("纯音频路径：" + outputFile.getAbsolutePath());
                    ToastUtils.showLong("纯音频路径：" + outputFile.getAbsolutePath());
                    LogUtils.i("extractorAndMuxerMP4: channelCount=" + trackFormat.getInteger(MediaFormat.KEY_CHANNEL_COUNT) + " bitRate=" + trackFormat.getInteger(MediaFormat.KEY_BIT_RATE));
                }
            }
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            mediaExtractor.release();
        }
    }

    /**
     * @param context
     * @param mediaExtractor
     * @param i
     * @param outputName
     * @return
     * @throws IOException
     */
    private File extractorAndMuxer(Context context, MediaExtractor mediaExtractor, int i, String outputName) throws IOException {
        MediaFormat trackFormat = mediaExtractor.getTrackFormat(i);
        MediaMuxer mediaMuxer;
        mediaExtractor.selectTrack(i);

        File outputFile = new File(context.getExternalCacheDir().getAbsolutePath() + "/MediaRecord" + outputName);
        if (outputFile.exists()) {
            outputFile.delete();
        }
        //1. 构造MediaMuxer
        mediaMuxer = new MediaMuxer(outputFile.getAbsolutePath(), MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4);
        //2. 添加轨道信息 参数为MediaFormat
        mediaMuxer.addTrack(trackFormat);
        //3. 开始合成
        mediaMuxer.start();

        //4. 设置buffer
        ByteBuffer buffer = ByteBuffer.allocate(500 * 1024);
        MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();

        //5.通过mediaExtractor.readSampleData读取数据流
        int sampleSize = 0;
        while ((sampleSize = mediaExtractor.readSampleData(buffer, 0)) > 0) {
            bufferInfo.flags = mediaExtractor.getSampleFlags();
            bufferInfo.offset = 0;
            bufferInfo.size = sampleSize;
            bufferInfo.presentationTimeUs = mediaExtractor.getSampleTime();
            int isEOS = bufferInfo.flags & MediaCodec.BUFFER_FLAG_END_OF_STREAM;
            LogUtils.i("demuxerMp4:  flags=" + bufferInfo.flags + " size=" + sampleSize + " time=" + bufferInfo.presentationTimeUs + " outputName" + outputName + " isEOS=" + isEOS);
            //6. 把通过mediaExtractor解封装的数据通过writeSampleData写入到对应的轨道
            mediaMuxer.writeSampleData(0, buffer, bufferInfo);
            mediaExtractor.advance();
        }
        LogUtils.i("extractorAndMuxer: " + outputName + "提取封装完成");

        mediaExtractor.unselectTrack(i);
        //6.关闭
        mediaMuxer.stop();
        mediaMuxer.release();
        return outputFile;
    }

    /**
     * 把音轨和视频轨再合成新的视频
     */
    public String muxerMp4() {
        File inputVideoFile = new File(mContext.getExternalCacheDir().getAbsolutePath() + "/MediaRecord/single_video.mp4");
        File inputAudioFile = new File(mContext.getExternalCacheDir().getAbsolutePath() + "/MediaRecord/single_audio.aac");
        File outputFile = new File(mContext.getExternalCacheDir().getAbsolutePath() + "/MediaRecord/muxer_video_record.mp4");

        if (outputFile.exists()) {
            outputFile.delete();
        }
        if (!inputVideoFile.exists()) {
            Toast.makeText(mContext, "视频源文件不存在", Toast.LENGTH_SHORT).show();
            return "";
        }
        if (!inputAudioFile.exists()) {
            Toast.makeText(mContext, "音频源文件不存在", Toast.LENGTH_SHORT).show();
            return "";
        }

        MediaExtractor videoExtractor = new MediaExtractor();
        MediaExtractor audioExtractor = new MediaExtractor();

        try {
            MediaMuxer mediaMuxer = new MediaMuxer(outputFile.getAbsolutePath(), MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4);
            int videoTrackIndex = 0;
            int audioTrackIndex = 0;

            //先添加视频轨道
            videoExtractor.setDataSource(inputVideoFile.getAbsolutePath());
            int trackCount = videoExtractor.getTrackCount();
            LogUtils.i("muxerToMp4: trackVideoCount=" + trackCount);

            for (int i = 0; i < trackCount; i++) {
                MediaFormat trackFormat = videoExtractor.getTrackFormat(i);
                String mimeType = trackFormat.getString(MediaFormat.KEY_MIME);
                if (TextUtils.isEmpty(mimeType)) {
                    continue;
                }
                if (mimeType.startsWith("video/")) {
                    videoExtractor.selectTrack(i);

                    videoTrackIndex = mediaMuxer.addTrack(trackFormat);
                    LogUtils.i("muxerToMp4: videoTrackIndex=" + videoTrackIndex);
                    break;
                }
            }

            //再添加音频轨道
            audioExtractor.setDataSource(inputAudioFile.getAbsolutePath());
            int trackCountAduio = audioExtractor.getTrackCount();
            LogUtils.i("muxerToMp4: trackCountAduio=" + trackCountAduio);
            for (int i = 0; i < trackCountAduio; i++) {
                MediaFormat trackFormat = audioExtractor.getTrackFormat(i);
                String mimeType = trackFormat.getString(MediaFormat.KEY_MIME);
                if (TextUtils.isEmpty(mimeType)) {
                    continue;
                }
                if (mimeType.startsWith("audio/")) {
                    audioExtractor.selectTrack(i);
                    audioTrackIndex = mediaMuxer.addTrack(trackFormat);
                    LogUtils.i("muxerToMp4: audioTrackIndex=" + audioTrackIndex);
                    break;
                }
            }

            //再进行合成
            mediaMuxer.start();

            ByteBuffer byteBuffer = ByteBuffer.allocate(500 * 1024);

            MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
            int sampleSize = 0;

            while ((sampleSize = videoExtractor.readSampleData(byteBuffer, 0)) > 0) {
                bufferInfo.flags = videoExtractor.getSampleFlags();
                bufferInfo.offset = 0;
                bufferInfo.size = sampleSize;
                bufferInfo.presentationTimeUs = videoExtractor.getSampleTime();
                mediaMuxer.writeSampleData(videoTrackIndex, byteBuffer, bufferInfo);
                videoExtractor.advance();
            }

            int audioSampleSize = 0;

            MediaCodec.BufferInfo audioBufferInfo = new MediaCodec.BufferInfo();

            while ((audioSampleSize = audioExtractor.readSampleData(byteBuffer, 0)) > 0) {
                audioBufferInfo.flags = audioExtractor.getSampleFlags();
                audioBufferInfo.offset = 0;
                audioBufferInfo.size = audioSampleSize;
                audioBufferInfo.presentationTimeUs = audioExtractor.getSampleTime();
                mediaMuxer.writeSampleData(audioTrackIndex, byteBuffer, audioBufferInfo);
                audioExtractor.advance();
            }

            //最后释放资源
            videoExtractor.release();
            audioExtractor.release();
            mediaMuxer.stop();
            mediaMuxer.release();

        } catch (IOException e) {
            e.printStackTrace();
            return "";

        }
        return outputFile.getAbsolutePath();
    }

}
